package com.kingdomai.testdb;

import com.alibaba.druid.sql.SQLUtils;
import com.alibaba.druid.sql.ast.SQLExpr;
import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr;
import com.alibaba.druid.sql.ast.expr.SQLValuableExpr;
import com.alibaba.druid.sql.ast.statement.*;
import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlSelectQueryBlock;
import com.alibaba.druid.util.JdbcConstants;

import java.util.*;
import java.util.stream.Collectors;

/**
 * @author yangzhixiao
 */
public class Database {

    public static Database INSTANCE = new Database();

    private final HashMap<String, DataTable> tableHashMap = new HashMap<>();

    public DataTable getTable(String name) {
        return tableHashMap.get(name);
    }

    /**
     * 解析SQL
     * @param sql
     */
    public void parseSQL(String sql) {
        List<SQLStatement> statements = SQLUtils.parseStatements(sql, JdbcConstants.MYSQL);
        for (SQLStatement statement : statements) {
            if (statement instanceof SQLCreateTableStatement) {
                // 建表
                createTable((SQLCreateTableStatement) statement);
            } else if (statement instanceof SQLDropTableStatement) {
                // 删表
                dropTable((SQLDropTableStatement) statement);
            } else if (statement instanceof SQLInsertStatement) {
                // 插入
                insert((SQLInsertStatement) statement);
            } else if (statement instanceof SQLSelectStatement) {
                // 查询
                query((SQLSelectStatement) statement);
            }
        }
    }

    /**
     * 查询语句
     * @param statement
     */
    private void query(SQLSelectStatement statement) {
        // 处理表关联
        MySqlSelectQueryBlock query = (MySqlSelectQueryBlock) statement.getSelect().getQuery();
        // 别名表数据
        Map<String, DataTable> aliasMap = new HashMap<>();
        processJoin(aliasMap, query.getFrom());
        SQLTableSource from = query.getFrom();
        while (from != null) {
            if (from instanceof SQLJoinTableSource) {
                SQLJoinTableSource jts = (SQLJoinTableSource) from;
                from = jts.getLeft();
            } else if (from instanceof SQLExprTableSource) {
                SQLExprTableSource ets = (SQLExprTableSource) from;
                from = null;
            } else {
                from = null;
            }
        }
    }

    private void processJoin(Map<String, DataTable> aliasMap, SQLTableSource left) {

        if (left instanceof SQLJoinTableSource && ((SQLJoinTableSource) left).getLeft() != null) {
            processJoin(aliasMap, ((SQLJoinTableSource) left).getLeft());
        }
    }

    /**
     * 建表语句
     * @param statement
     */
    public void createTable(SQLCreateTableStatement statement) {
        String tableName = Utils.getSimpleName(statement.getName());
        DataTable table = new DataTable();
        table.setTableName(tableName);
        List<SQLColumnDefinition> columnDefinitions = new LinkedList<>();
        for (SQLTableElement element : statement.getTableElementList()) {
            if (element instanceof SQLPrimaryKey) {
                table.setPrimaryKey((SQLPrimaryKey) element);
            } else if (element instanceof SQLColumnDefinition) {
                columnDefinitions.add((SQLColumnDefinition) element);
            }
        }
        table.setColumnDefinitions(columnDefinitions);
        tableHashMap.put(tableName, table);
    }

    /**
     * 删表
     * @param statement
     */
    public void dropTable(SQLDropTableStatement statement) {
        for (SQLExprTableSource tableSource : statement.getTableSources()) {
            String tableName = Utils.getSimpleName(tableSource.getName());
            if (tableHashMap.containsKey(tableName)) {
                tableHashMap.remove(tableName);
            } else {
                throw new TDException(String.format("Table [%s] does not existed.", tableName));
            }
        }
    }

    /**
     * 插入记录
     * @param statement
     */
    public void insert(SQLInsertStatement statement) {
        String tableName = Utils.getSimpleName(statement.getTableName());
        DataTable table = tableHashMap.get(tableName);
        if (table == null) {
            throw new TDException(String.format("Table [%s] does not existed.", tableName));
        }
        List<SQLExpr> columns = statement.getColumns();
        for (SQLInsertStatement.ValuesClause valuesClause : statement.getValuesList()) {
            List<SQLExpr> values = valuesClause.getValues();
            Map<String, Object> row = getInsertValues(
                    table,
                    columns.stream().map(c -> (SQLIdentifierExpr) c).collect(Collectors.toList()),
                    values.stream().map(v -> (SQLValuableExpr) v).collect(Collectors.toList())
            );
            String pkValue = getPkStrValue(table.getPrimaryKey(), row);
            // 主键duplicate key约束检查
            if (table.getRows().containsKey(pkValue)) {
                throw new TDException(String.format("Duplicate key for [%s]", pkValue));
            }
            table.getRows().put(pkValue, row);
        }
    }

    public Map<String, Object> getInsertValues(DataTable table, List<SQLIdentifierExpr> insertColumns, List<SQLValuableExpr> insertValues) {
        Map<String, Object> row = new HashMap<>();
        for (SQLColumnDefinition columnDefinition : table.getColumnDefinitions()) {
            Object value = null;
            // 默认值
            if (!Objects.isNull(columnDefinition.getDefaultExpr())) {
                // TODO 表达式优化
                value = columnDefinition.getDefaultExpr().toString();
            }
            // 取列值
            for (int i = 0; i < insertColumns.size(); i++) {
                if (insertColumns.get(i).nameEquals(columnDefinition.getColumnName())) {
                    value = insertValues.get(i).getValue();
                    break;
                }
            }
            // 列值空约束检查
            if (columnDefinition.containsNotNullConstaint() && Objects.isNull(value)) {
                throw new TDException(String.format("Constaint failed not null for %s", columnDefinition.toString()));
            }
            // 主键空约束检查
            if (columnDefinition.isPrimaryKey() && Objects.isNull(value)) {
                throw new TDException(String.format("Primary key can not be null for %s", columnDefinition.toString()));
            }
            row.put(Utils.getSimpleName(columnDefinition.getName()), value);
        }
        return row;
    }

    public String getPkStrValue(SQLPrimaryKey primaryKey, Map<String, Object> row) {
        List<String> values = new ArrayList<>();
        for (SQLSelectOrderByItem column : primaryKey.getColumns()) {
            values.add(row.get(column.toString().replaceAll("`", "")).toString());
        }
        return String.join("-", values);
    }

}
